from collections import OrderedDict
from os import path as osp
from tqdm import tqdm

import torch
import torchvision.utils as tvu
import copy

from basicsr.archs import build_network
from basicsr.losses import build_loss
from basicsr.utils import get_root_logger, imwrite, tensor2img, img2tensor
from basicsr.utils.registry import MODEL_REGISTRY
from .base_model import BaseModel
import pyiqa


@MODEL_REGISTRY.register()
class DualCodeBookModel(BaseModel):
    """
    Model class for DualCodeBookNet.

    It handles dual inputs (hazy and clear), dual outputs, computes losses for both branches,
    and employs two separate discriminators for adversarial training.
    """

    def __init__(self, opt):
        super().__init__(opt)
        # define network
        self.net_g = build_network(opt['network_g'])
        self.net_g = self.model_to_device(self.net_g)
        self.show_codebook_usage = opt['network_g'].get('show_codebook_usage', False)
        self.LQ_stage = self.opt['network_g'].get('LQ_stage', False)
        self.quantize_way = self.opt['network_g'].get('quantize_way', None)
        self.show_dist_norm = self.opt['network_g'].get('show_dist_norm', False)
        if self.quantize_way == None:
            raise ValueError("Please specify 'quantize_way' in network_g options.")
        if self.LQ_stage:
            self.former = self.opt['network_g'].get('former', False)
        # define metric functions
        if self.opt['val'].get('metrics') is not None:
            self.metric_funcs = {}
            for _, opt in self.opt['val']['metrics'].items():
                mopt = opt.copy()
                name = mopt.pop('type', None)
                mopt.pop('better', None)
                self.metric_funcs[name] = pyiqa.create_metric(name, device=self.device, **mopt)

        # load pre-trained HQ ckpt, frozen decoder and codebook
        
        if self.LQ_stage:
            load_path = self.opt['path'].get('pretrain_network_hq', None)
            assert load_path is not None, 'Need to specify hq prior model path in LQ stage'
            self.load_network(self.net_g, load_path, False)
            frozen_module_keywords = self.opt['network_g'].get('frozen_module_keywords', None)
            if frozen_module_keywords is not None:
                for name, module in self.net_g.named_modules():
                    for fkw in frozen_module_keywords:
                        if fkw in name:
                            for p in module.parameters():
                                p.requires_grad = False
                            break
        else:
            load_path_g = self.opt['path'].get('pretrain_network_g', None)
            logger = get_root_logger()
            if load_path_g is not None:
                logger.info(f'Loading net_g from {load_path_g}')
                # self.load_network(self.net_g, load_path_g, self.opt['path']['strict_load'])
                self.load_femasr_into_dualcodebook(self.net_g, load_path_g)


        if self.is_train:
            self.init_training_settings()
            # 【逻辑保留】依然保留gan_opt来控制是否使用判别器
            self.use_dis = (self.opt['train']['gan_opt']['loss_weight'] != 0)
            
            # 【代码差异】为hazy和clear两路输出分别创建最佳判别器的副本
            if not self.LQ_stage:
                self.net_d_hazy_best = copy.deepcopy(self.net_d_hazy)
            self.net_d_clear_best = copy.deepcopy(self.net_d_clear)
    

        self.net_g_best = copy.deepcopy(self.net_g)

    def init_training_settings(self):
        logger = get_root_logger()
        train_opt = self.opt['train']
        self.net_g.train()

        # 【代码差异】分别加载两个判别器的预训练模型
        load_path_d_hazy = self.opt['path'].get('pretrain_network_d_hazy', None)
        if load_path_d_hazy is not None:
            self.net_d_hazy = build_network(self.opt['network_d_hazy'])
            self.net_d_hazy = self.model_to_device(self.net_d_hazy)
            logger.info(f'Loading net_d_hazy from {load_path_d_hazy}')
            self.load_network(self.net_d_hazy, load_path_d_hazy, self.opt['path'].get('strict_load_d', True))
            self.net_d_hazy.train()
        
        load_path_d_clear = self.opt['path'].get('pretrain_network_d_clear', None)
        if load_path_d_clear is not None:
            self.net_d_clear = build_network(self.opt['network_d_clear'])
            self.net_d_clear = self.model_to_device(self.net_d_clear)

            logger.info(f'Loading net_d_clear from {load_path_d_clear}')
            self.load_network(self.net_d_clear, load_path_d_clear, self.opt['path'].get('strict_load_d', True))
            self.net_d_clear.train()


        # define losses (与FeMaSRModel共享相同的损失函数定义)
        if train_opt.get('pixel_opt'):
            self.cri_pix = build_loss(train_opt['pixel_opt']).to(self.device)
        else:
            self.cri_pix = None

        if train_opt.get('perceptual_opt'):
            self.cri_perceptual = build_loss(train_opt['perceptual_opt']).to(self.device)
            self.model_to_device(self.cri_perceptual)
        else:
            self.cri_perceptual = None

        if train_opt.get('gan_opt'):
            self.cri_gan = build_loss(train_opt['gan_opt']).to(self.device)

        self.net_d_iters = train_opt.get('net_d_iters', 1)
        self.net_d_init_iters = train_opt.get('net_d_init_iters', 0)

        # set up optimizers and schedulers
        self.setup_optimizers()
        self.setup_schedulers()

    def setup_optimizers(self):
        train_opt = self.opt['train']
        
        # optimizer g (与FeMaSRModel逻辑相同)
        optim_params_g = []
        for k, v in self.net_g.named_parameters():
            if v.requires_grad:
                optim_params_g.append(v)
            else:
                logger = get_root_logger()
                logger.warning(f'Params {k} will not be optimized.')
        
        optim_type_g = train_opt['optim_g'].pop('type')
        self.optimizer_g = getattr(torch.optim, optim_type_g)(optim_params_g, **train_opt['optim_g'])
        self.optimizers.append(self.optimizer_g)

        # 【主要不同点】为两个判别器分别设置优化器
        # optimizer d for hazy branch
        if self.LQ_stage == False:
            optim_type_d_hazy = train_opt['optim_d_hazy']
            if optim_type_d_hazy != None:
                optim_type_d_hazy = optim_type_d_hazy.pop('type')
                self.optimizer_d_hazy = getattr(torch.optim, optim_type_d_hazy)(self.net_d_hazy.parameters(), **train_opt['optim_d_hazy'])
                self.optimizers.append(self.optimizer_d_hazy)
        
        # optimizer d for clear branch
        optim_type_d_clear = train_opt['optim_d_clear']
        if optim_type_d_clear != None:
            optim_type_d_clear = optim_type_d_clear.pop('type')
            self.optimizer_d_clear = getattr(torch.optim, optim_type_d_clear)(self.net_d_clear.parameters(), **train_opt['optim_d_clear'])
            self.optimizers.append(self.optimizer_d_clear)


    def feed_data(self, data):
        self.gt = data['gt'].to(self.device)
        self.lq = data['lq'].to(self.device)  # Hazy images

    def optimize_parameters(self, current_iter):
        train_opt = self.opt['train']

        # ----------------- Optimize Generator (net_g) ----------------- #
        # 冻结两个判别器
        if self.LQ_stage == False:
            for p in self.net_d_hazy.parameters():
                p.requires_grad = False
            for p in self.net_d_clear.parameters():
                p.requires_grad = False
        else:
            for p in self.net_d_clear.parameters():
                p.requires_grad = False
        
        
        self.optimizer_g.zero_grad()
        if self.LQ_stage == False:
            self.output_hazy, self.output_clear, l_codebook, l_semantic, self.codebook_usage, (self.d_hazy_norm, self.d_clear_norm) = self.net_g(self.lq, self.gt)
        else:
            self.output_hazy, self.output_clear, l_codebook, l_semantic, self.codebook_usage, (self.d_hazy_norm, self.d_clear_norm) = self.net_g(self.lq, None)
        l_g_total = 0
        loss_dict = OrderedDict()
        
        # codebook loss (与FeMaSRModel逻辑相同，由网络内部返回)
        if train_opt.get('codebook_opt', None):
            l_codebook *= train_opt['codebook_opt']['loss_weight']
            l_g_total += l_codebook.mean()
            loss_dict['l_codebook'] = l_codebook.mean()

        # semantic loss
        if train_opt.get('semantic_opt', None) and isinstance(l_semantic, torch.Tensor):
            l_semantic *= train_opt['semantic_opt']['loss_weight']
            l_g_total += l_semantic.mean()
            loss_dict['l_semantic'] = l_semantic.mean()

        # 【主要不同点】计算两部分的重建损失
        # pixel loss for both hazy and clear branches
        if self.cri_pix:
            if self.LQ_stage == False:
                l_pix_hazy = self.cri_pix(self.output_hazy, self.lq) # hazy输出与hazy输入比较
                l_pix_clear = self.cri_pix(self.output_clear, self.gt) # clear输出与clear输入比较
                l_pix = l_pix_hazy + l_pix_clear
                l_g_total += l_pix
                loss_dict['l_pix_hazy'] = l_pix_hazy
                loss_dict['l_pix_clear'] = l_pix_clear
            else:
                l_pix_clear = self.cri_pix(self.output_clear, self.gt) # clear输出与clear输入比较
                l_g_total += l_pix_clear
                loss_dict['l_pix_clear'] = l_pix_clear

        # perceptual loss for both hazy and clear branches
        if self.cri_perceptual:
            if self.LQ_stage == False:
                # Hazy branch
                l_percep_hazy, l_style_hazy = self.cri_perceptual(self.output_hazy, self.lq)
                if l_percep_hazy is not None:
                    l_g_total += l_percep_hazy.mean()
                    loss_dict['l_percep_hazy'] = l_percep_hazy.mean()
                if l_style_hazy is not None:
                    l_g_total += l_style_hazy
                    loss_dict['l_style_hazy'] = l_style_hazy
                # Clear branch
                l_percep_clear, l_style_clear = self.cri_perceptual(self.output_clear, self.gt)
                if l_percep_clear is not None:
                    l_g_total += l_percep_clear.mean()
                    loss_dict['l_percep_clear'] = l_percep_clear.mean()
                if l_style_clear is not None:
                    l_g_total += l_style_clear
                    loss_dict['l_style_clear'] = l_style_clear
            else:
                # Clear branch only
                l_percep_clear, l_style_clear = self.cri_perceptual(self.output_clear, self.gt)
                if l_percep_clear is not None:
                    l_g_total += l_percep_clear.mean()
                    loss_dict['l_percep_clear'] = l_percep_clear.mean()
                if l_style_clear is not None:
                    l_g_total += l_style_clear
                    loss_dict['l_style_clear'] = l_style_clear

        # 【主要不同点】计算两部分的GAN损失
        # gan loss for both hazy and clear branches
        if self.use_dis and current_iter > train_opt['net_d_init_iters']:
            if self.LQ_stage == False:
                fake_g_pred_hazy = self.net_d_hazy(self.output_hazy)
                l_g_gan_hazy = self.cri_gan(fake_g_pred_hazy, True, is_disc=False)
                
                fake_g_pred_clear = self.net_d_clear(self.output_clear)
                l_g_gan_clear = self.cri_gan(fake_g_pred_clear, True, is_disc=False)

                l_g_gan = l_g_gan_hazy + l_g_gan_clear
                l_g_total += l_g_gan
                loss_dict['l_g_gan_hazy'] = l_g_gan_hazy
                loss_dict['l_g_gan_clear'] = l_g_gan_clear
            else:
                fake_g_pred_clear = self.net_d_clear(self.output_clear)
                l_g_gan_clear = self.cri_gan(fake_g_pred_clear, True, is_disc=False)
                l_g_total += l_g_gan_clear
                loss_dict['l_g_gan_clear'] = l_g_gan_clear

        l_g_total.mean().backward()
        self.optimizer_g.step()

        # ----------------- Optimize Discriminators ----------------- #
        if self.use_dis and current_iter > train_opt['net_d_init_iters']:
            if self.LQ_stage == False:
                # Optimize net_d_hazy
                for p in self.net_d_hazy.parameters():
                    p.requires_grad = True
                self.optimizer_d_hazy.zero_grad()
                # real (for net_d_hazy, "real" means a real hazy image)
                real_d_pred_hazy = self.net_d_hazy(self.lq)
                l_d_real_hazy = self.cri_gan(real_d_pred_hazy, True, is_disc=True)
                loss_dict['l_d_real_hazy'] = l_d_real_hazy
                loss_dict['out_d_real_hazy'] = torch.mean(real_d_pred_hazy.detach())
                l_d_real_hazy.backward()
                # fake
                fake_d_pred_hazy = self.net_d_hazy(self.output_hazy.detach())
                l_d_fake_hazy = self.cri_gan(fake_d_pred_hazy, False, is_disc=True)
                loss_dict['l_d_fake_hazy'] = l_d_fake_hazy
                loss_dict['out_d_fake_hazy'] = torch.mean(fake_d_pred_hazy.detach())
                l_d_fake_hazy.backward()
                self.optimizer_d_hazy.step()

                # Optimize net_d_clear
                for p in self.net_d_clear.parameters():
                    p.requires_grad = True
                self.optimizer_d_clear.zero_grad()
                # real (for net_d_clear, "real" means a real clear image)
                real_d_pred_clear = self.net_d_clear(self.gt)
                l_d_real_clear = self.cri_gan(real_d_pred_clear, True, is_disc=True)
                loss_dict['l_d_real_clear'] = l_d_real_clear
                loss_dict['out_d_real_clear'] = torch.mean(real_d_pred_clear.detach())
                l_d_real_clear.backward()
                # fake
                fake_d_pred_clear = self.net_d_clear(self.output_clear.detach())
                l_d_fake_clear = self.cri_gan(fake_d_pred_clear, False, is_disc=True)
                loss_dict['l_d_fake_clear'] = l_d_fake_clear
                loss_dict['out_d_fake_clear'] = torch.mean(fake_d_pred_clear.detach())
                l_d_fake_clear.backward()
                self.optimizer_d_clear.step()
            else:
                # Optimize net_d_clear only
                for p in self.net_d_clear.parameters():
                    p.requires_grad = True
                self.optimizer_d_clear.zero_grad()
                # real
                real_d_pred_clear = self.net_d_clear(self.gt)
                l_d_real_clear = self.cri_gan(real_d_pred_clear, True, is_disc=True)
                loss_dict['l_d_real_clear'] = l_d_real_clear
                loss_dict['out_d_real_clear'] = torch.mean(real_d_pred_clear.detach())
                l_d_real_clear.backward()
                # fake
                fake_d_pred_clear = self.net_d_clear(self.output_clear.detach())
                l_d_fake_clear = self.cri_gan(fake_d_pred_clear, False, is_disc=True)
                loss_dict['l_d_fake_clear'] = l_d_fake_clear
                loss_dict['out_d_fake_clear'] = torch.mean(fake_d_pred_clear.detach())
                l_d_fake_clear.backward()
                self.optimizer_d_clear.step()

        self.log_dict = self.reduce_loss_dict(loss_dict)

    def reset_usage(self):
        """在每个 epoch 开始时调用，重置使用情况统计。"""
        if self.show_codebook_usage:
            net_g = self.get_bare_model(self.net_g)
            net_g.quantize_group[0].reset_usage()

    
    def test(self):
        """
        [修改] test函数现在只负责执行前向传播，不再做选择性赋值。
        它将两个输出都保留在实例变量中，供调用者（如validation函数）使用。
        """
        self.net_g.eval()
        net_g = self.get_bare_model(self.net_g)
        
        with torch.no_grad():
            if self.LQ_stage == False:
                self.output_hazy, self.output_clear, _, _, self.codebook_usage, (self.d_hazy_norm, self.d_clear_norm)= net_g(self.lq, self.gt)
            else:
                self.output_hazy, self.output_clear, _, _, self.codebook_usage, (self.d_hazy_norm, self.d_clear_norm)= net_g(self.lq, None)
        
        self.net_g.train()

    def dist_validation(self, dataloader, current_iter, tb_logger, save_img, save_as_dir=None):
        logger = get_root_logger()
        logger.info('Only support single GPU validation.')
        self.nondist_validation(dataloader, current_iter, tb_logger, save_img, save_as_dir)

    def nondist_validation(self, dataloader, current_iter, tb_logger,
                           save_img, save_as_dir):
        """
        [重大修改] 此函数被重构以分别评估hazy和clear两个分支的重建性能。
        """
        dataset_name = dataloader.dataset.opt['name']
        with_metrics = self.opt['val'].get('metrics') is not None
        
        if with_metrics:
            if self.LQ_stage == False:
                # 【修改】为hazy和clear分支分别初始化指标结果字典
                self.metric_results_hazy = {metric: 0 for metric in self.opt['val']['metrics'].keys()}
                self.metric_results_clear = {metric: 0 for metric in self.opt['val']['metrics'].keys()}
                
                # 这里的best_metric_results也需要能区分hazy和clear，或者我们只用一个关键指标来保存模型
                # 为简单起见，我们假设clear分支的指标是保存模型的关键
                self._initialize_best_metric_results(dataset_name) 
                self.key_metric = self.opt['val'].get('key_metric') 
            else:
                self.metric_results_clear = {metric: 0 for metric in self.opt['val']['metrics'].keys()}
                self._initialize_best_metric_results(dataset_name) 
                self.key_metric = self.opt['val'].get('key_metric')
        pbar = tqdm(total=len(dataloader), unit='image')
        if self.LQ_stage == False:
            for idx, val_data in enumerate(dataloader):
                img_name = osp.splitext(osp.basename(val_data['lq_path'][0]))[0]
                self.feed_data(val_data)
                self.test()
                # --- 分别处理 Hazy 和 Clear 分支 ---
                # 1. Hazy 分支处理
                img_hazy_out = tensor2img(self.output_hazy)
                if save_img:
                    save_img_path_hazy = osp.join(self.opt['path']['visualization'], dataset_name, f'{img_name}_hazy_out.png')
                    imwrite(img_hazy_out, save_img_path_hazy)
                if with_metrics:
                    metric_data_hazy = [img2tensor(img_hazy_out).unsqueeze(0) / 255., self.lq]
                    for name, opt_ in self.opt['val']['metrics'].items():
                        self.metric_results_hazy[name] += self.metric_funcs[name](*metric_data_hazy).item()
                # 2. Clear 分支处理
                img_clear_out = tensor2img(self.output_clear)
                if save_img:
                    save_img_path_clear = osp.join(self.opt['path']['visualization'], dataset_name, f'{img_name}_clear_out.png')
                    imwrite(img_clear_out, save_img_path_clear)
                if with_metrics:
                    metric_data_clear = [img2tensor(img_clear_out).unsqueeze(0) / 255., self.gt]
                    for name, opt_ in self.opt['val']['metrics'].items():
                        self.metric_results_clear[name] += self.metric_funcs[name](*metric_data_clear).item()
                # 释放显存
                del self.lq, self.gt, self.output_hazy, self.output_clear
                
                pbar.update(1)
                pbar.set_description(f'Test {img_name}')
        else:
            for idx, val_data in enumerate(dataloader):
                img_name = osp.splitext(osp.basename(val_data['lq_path'][0]))[0]
                self.feed_data(val_data)
                self.test()
                # --- 仅处理 Clear 分支 ---
                img_clear_out = tensor2img(self.output_clear)
                if save_img:
                    save_img_path_clear = osp.join(self.opt['path']['visualization'], dataset_name, f'{img_name}_clear_out.png')
                    imwrite(img_clear_out, save_img_path_clear)
                if with_metrics:
                    metric_data_clear = [img2tensor(img_clear_out).unsqueeze(0) / 255., self.gt]
                    for name, opt_ in self.opt['val']['metrics'].items():
                        self.metric_results_clear[name] += self.metric_funcs[name](*metric_data_clear).item()
                # 释放显存
                del self.lq, self.gt, self.output_clear
                
                pbar.update(1)
                pbar.set_description(f'Test {img_name}')

        pbar.close()
        if self.LQ_stage == False:
            if with_metrics:
                # 【修改】计算两个分支的平均指标
                for metric in self.metric_results_hazy.keys():
                    self.metric_results_hazy[metric] /= (idx + 1)
                    self.metric_results_clear[metric] /= (idx + 1)       
                # 【修改】日志记录和模型保存逻辑
                # 我们仍然使用clear分支的指标作为保存"best"模型的依据
                if self.key_metric is not None:
                    if self._update_best_metric_result(dataset_name, self.key_metric, self.metric_results_clear[self.key_metric], current_iter):
                        self.copy_model(self.net_g, self.net_g_best)
                        self.copy_model(self.net_d_hazy, self.net_d_hazy_best)
                        self.copy_model(self.net_d_clear, self.net_d_clear_best)
                        self.save_network(self.net_g, 'net_g_best', '')
                        self.save_network(self.net_d_hazy, 'net_d_hazy_best', '')
                        self.save_network(self.net_d_clear, 'net_d_clear_best', '')
                
                # 记录两个分支的指标到日志和TensorBoard
                self._log_validation_metric_values(current_iter, dataset_name, tb_logger)
        else:
            if with_metrics:
                # 【修改】计算两个分支的平均指标
                for metric in self.metric_results_clear.keys():
                    self.metric_results_clear[metric] /= (idx + 1)       
                # 【修改】日志记录和模型保存逻辑
                # 我们仍然使用clear分支的指标作为保存"best"模型的依据
                if self.key_metric is not None:
                    if self._update_best_metric_result(dataset_name, self.key_metric, self.metric_results_clear[self.key_metric], current_iter):
                        self.copy_model(self.net_g, self.net_g_best)
                        self.copy_model(self.net_d_clear, self.net_d_clear_best)
                        self.save_network(self.net_g, 'net_g_best', '')
                        self.save_network(self.net_d_clear, 'net_d_clear_best', '')
                
                # 记录两个分支的指标到日志和TensorBoard
                self._log_validation_metric_values(current_iter, dataset_name, tb_logger)


    def _log_validation_metric_values(self, current_iter, dataset_name, tb_logger):
        """
        [重大修改] 日志函数现在需要记录 Hazy 和 Clear 两个分支的指标.
        """
        logger = get_root_logger()
        
        # Log for Clear branch
        log_str_clear = f'Validation {dataset_name} - Clear Reconstruction\n'
        for metric, value in self.metric_results_clear.items():
            log_str_clear += f'\t # Clear_{metric}: {value:.4f}'
            # 最佳指标仍然显示clear分支的
            if hasattr(self, 'best_metric_results') and metric in self.best_metric_results[dataset_name]:
                log_str_clear += (f'\tBest: {self.best_metric_results[dataset_name][metric]["val"]:.4f} @ '
                                  f'{self.best_metric_results[dataset_name][metric]["iter"]} iter')
            log_str_clear += '\n'
        logger.info(log_str_clear)


        if tb_logger:
            for metric, value in self.metric_results_clear.items():
                tb_logger.add_scalar(f'metrics/{dataset_name}/Clear_{metric}', value, current_iter)
    
    def vis_single_code(self, up_factor=4):
        # 【逻辑相同】可视化码本的功能可以保留
        net_g = self.get_bare_model(self.net_g)
        code_num = self.opt['network_g']['codebook_params'][0][1] 
        with torch.no_grad():
            code_idx = torch.arange(code_num).reshape(code_num, 1, 1, 1) # (code_num, 1, 1, 1)
            code_idx = code_idx.repeat(1, 1, up_factor, up_factor) # (code_num, 1, 4, 4)
            output_img = net_g.decode_indices(code_idx) 
            output_img_hazy = tvu.make_grid(output_img[0], nrow=32)
            output_img_clear = tvu.make_grid(output_img[1], nrow=32)

        return output_img_hazy.unsqueeze(0), output_img_clear.unsqueeze(0)

    def get_current_visuals(self):
        # 【主要不同点】可视化部分需要展示更多的图像
        if not self.LQ_stage:
            vis_samples = 16
            out_dict = OrderedDict()
            out_dict['lq_hazy'] = self.lq.detach().cpu()[:vis_samples]
            out_dict['gt_clear'] = self.gt.detach().cpu()[:vis_samples]
            out_dict['output_hazy'] = self.output_hazy.detach().cpu()[:vis_samples]
            out_dict['output_clear'] = self.output_clear.detach().cpu()[:vis_samples]
            net_g_bare = self.get_bare_model(self.net_g)
            if not self.LQ_stage and hasattr(net_g_bare, 'decode_indices'):
                out_dict['codebook_hazy'], out_dict['codebook_clear'] = self.vis_single_code()
        else:
            vis_samples = 16
            out_dict = OrderedDict()
            out_dict['lq_hazy'] = self.lq.detach().cpu()[:vis_samples]
            out_dict['output_clear'] = self.output_clear.detach().cpu()[:vis_samples]
        
        return out_dict


    def get_codebook_usage(self):
        return self.codebook_usage
    
    def get_dist_norm(self):
        return self.d_hazy_norm, self.d_clear_norm

    def save(self, epoch, current_iter):
        # 【主要不同点】保存模型时需要保存一个生成器和两个判别器
        if self.LQ_stage == False:
            self.save_network(self.net_g, 'net_g', current_iter)
            self.save_network(self.net_d_hazy, 'net_d_hazy', current_iter)
            self.save_network(self.net_d_clear, 'net_d_clear', current_iter)
            self.save_training_state(epoch, current_iter)
        else:
            self.save_network(self.net_g, 'net_g', current_iter)
            self.save_network(self.net_d_clear, 'net_d_clear', current_iter)
            self.save_training_state(epoch, current_iter)

    def load_femasr_into_dualcodebook(self, model_dual, pretrained_path):
        print(f"Loading pretrained FeMaSRNet weights from: {pretrained_path}")
        
        # 1. 加载预训练权重
        load_net = torch.load(pretrained_path, map_location=lambda storage, loc: storage)
        
        if 'params_ema' in load_net:
            pretrained_state_dict = load_net['params_ema']
        elif 'params' in load_net:
            pretrained_state_dict = load_net['params']
        else:
            pretrained_state_dict = load_net
            
        # 2. 获取新模型 state_dict 及处理 'module.' 前缀
        new_model_state_dict = model_dual.state_dict()
        new_model_has_module_prefix = all(key.startswith('module.') for key in new_model_state_dict.keys())
        print(f"New model has 'module.' prefix: {new_model_has_module_prefix}")
        final_state_dict = OrderedDict()
        
        # 3. 遍历预训练权重并重映射
        mapped_count = 0
        skipped_keys = []
        
        for pretrained_key, value in pretrained_state_dict.items():
            key_without_prefix = pretrained_key[7:] if pretrained_key.startswith('module.') else pretrained_key
            target_keys_without_prefix = []
            is_mapped = False
            if key_without_prefix.startswith('multiscale_encoder.'):
                target_keys_without_prefix.append(key_without_prefix.replace('multiscale_encoder.', 'multiscale_encoder_hazy.', 1))
                target_keys_without_prefix.append(key_without_prefix.replace('multiscale_encoder.', 'multiscale_encoder_clear.', 1))
            
            elif key_without_prefix.startswith('decoder_group.'):
                target_keys_without_prefix.append(key_without_prefix.replace('decoder_group.', 'decoder_group_hazy.', 1))
                target_keys_without_prefix.append(key_without_prefix.replace('decoder_group.', 'decoder_group_clear.', 1))
            elif key_without_prefix.startswith('out_conv.'):
                target_keys_without_prefix.append(key_without_prefix.replace('out_conv.', 'out_conv_hazy.', 1))
                target_keys_without_prefix.append(key_without_prefix.replace('out_conv.', 'out_conv_clear.', 1))
            # --------------------------- 新增: 处理 quantize_group 的特殊情况 ---------------------------
            elif key_without_prefix == 'quantize_group.0.embedding.weight':
                print(f"--- Special custom handling for '{pretrained_key}' ---")
                
                # 确定最终的目标键名 (可能需要加 'module.')
                target_key_wo_prefix = key_without_prefix
                final_target_key = f'module.{target_key_wo_prefix}' if new_model_has_module_prefix else target_key_wo_prefix
                # 检查目标键是否存在于新模型
                if final_target_key in new_model_state_dict:
                    target_shape = new_model_state_dict[final_target_key].shape
                    source_shape = value.shape
                    
                    # 执行您需求的自定义逻辑：如果源和目标维度匹配要求
                    if len(source_shape) == 2 and len(target_shape) == 2 and \
                    source_shape[0] == target_shape[0] and \
                    source_shape[1] * 2 == target_shape[1]:
                        
                        print(f"Concatenating pre-trained weight of shape {source_shape} to match target shape {target_shape}.")
                        concatenated_value = torch.cat([value, value], dim=1)
                        
                        final_state_dict[final_target_key] = concatenated_value.clone()
                        mapped_count += 1
                        is_mapped = True
                    else:
                        # 如果维度不满足拼接条件，则按常规逻辑检查是否完全匹配
                        if source_shape == target_shape:
                            print(f"Shapes for '{pretrained_key}' already match. Loading directly.")
                            final_state_dict[final_target_key] = value.clone()
                            mapped_count += 1
                            is_mapped = True
                        else:
                            print(f"Warning: Shape mismatch for special key '{final_target_key}'. Pretrained: {source_shape}, New: {target_shape}. Skipping.")
                # 跳过.append()环节，直接进入下一次循环
                if not is_mapped: skipped_keys.append(pretrained_key)
                continue # 使用continue确保此特殊键不会被后续逻辑再次处理
            # ------------------------------------ 新增逻辑结束 -------------------------------------
            else: # 对于名称可能完全匹配的层
                target_keys_without_prefix.append(key_without_prefix)
            # 常规加载逻辑 (对于非特殊处理的键)
            for target_key_wo_prefix in target_keys_without_prefix:
                final_target_key = f'module.{target_key_wo_prefix}' if new_model_has_module_prefix else target_key_wo_prefix
                
                if final_target_key in new_model_state_dict:
                    if new_model_state_dict[final_target_key].shape == value.shape:
                        final_state_dict[final_target_key] = value.clone()
                        mapped_count += 1
                        is_mapped = True
                    else:
                        print(f"Warning: shapes for key '{final_target_key}' mismatch! Pretrained({pretrained_key}): {value.shape}, New: {new_model_state_dict[final_target_key].shape}. Skipping.")
            
            if not is_mapped:
                skipped_keys.append(pretrained_key)
        if skipped_keys:
            print("\nInfo: The following pretrained keys were not found or did not match in the new model:")
            for k in sorted(list(set(skipped_keys))): # 用set去重并排序
                print(f"- {k}")
        model_dual.load_state_dict(final_state_dict, strict=False)
        total_new_params = len(new_model_state_dict)
        print(f"\nSuccessfully loaded/adapted and mapped {len(final_state_dict)} tensors.")
        print(f"This corresponds to {len(final_state_dict)} / {total_new_params} ({len(final_state_dict)/total_new_params:.2%}) of the new model's parameters.")